#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
HiAgent Chat SDK 核心客户端模块

提供HiAgent API的Python封装，支持对话聊天、会话管理、工作流执行等功能。
"""

import json
import requests
import sseclient
from typing import Dict, List, Optional, Any, Iterator, Union
from urllib.parse import urljoin

from models import (
    Conversation, Message, ChatResponse, WorkflowResult, 
    LongMemory, UserVariable, QueryExtends, FileInfo
)
from exceptions import (
    HiAgentException, APIException, AuthenticationException,
    ValidationException, NetworkException, ConversationNotFoundException,
    MessageNotFoundException, WorkflowException, RateLimitException
)
from config import HiAgentConfig


class HiAgentClient:
    """HiAgent API客户端
    
    提供完整的HiAgent API调用功能，包括：
    - 对话会话管理
    - 流式和阻塞式聊天
    - 消息管理和反馈
    - 工作流执行
    - 长期记忆管理
    - 用户变量管理
    """
    
    def __init__(self, config: HiAgentConfig):
        """初始化客户端
        
        Args:
            config: HiAgent配置对象
        """
        self.config = config
        self.api_key = config.app_key
        self.base_url = config.base_url.rstrip('/')
        self.timeout = config.timeout
        self.session = requests.Session()
        self.session.headers.update({
            'Apikey': self.api_key,
            'Content-Type': 'application/json'
        })
    
    def _make_request(self, method: str, endpoint: str, data: dict = None, 
                     stream: bool = False, **kwargs) -> requests.Response:
        """发送HTTP请求
        
        Args:
            method: HTTP方法
            endpoint: API端点
            data: 请求数据
            stream: 是否流式请求
            **kwargs: 其他请求参数
            
        Returns:
            requests.Response: 响应对象
            
        Raises:
            NetworkException: 网络错误
            AuthenticationException: 认证错误
            APIException: API错误
        """
        url = urljoin(self.base_url, endpoint)
        
        try:
            if stream:
                kwargs['headers'] = kwargs.get('headers', {})
                kwargs['headers']['Accept'] = 'text/event-stream'
            
            response = self.session.request(
                method=method,
                url=url,
                json=data,
                timeout=self.timeout,
                stream=stream,
                **kwargs
            )
            
            # 检查认证错误
            if response.status_code == 401:
                raise AuthenticationException("API Key无效或已过期")
            
            # 检查频率限制
            if response.status_code == 429:
                raise RateLimitException("请求频率超过限制")
            
            # 检查其他HTTP错误
            if not stream and not response.ok:
                try:
                    error_data = response.json()
                    error_msg = error_data.get('message', f'HTTP {response.status_code}')
                except:
                    error_msg = f'HTTP {response.status_code}: {response.text}'

                # Log详细的错误信息用于调试
                import logging
                logger = logging.getLogger(__name__)
                logger.error(f"API request failed: {method} {url}")
                logger.error(f"Status: {response.status_code}")
                logger.error(f"Response: {response.text}")
                logger.error(f"Request data: {data}")

                raise APIException(
                    message=error_msg,
                    status_code=response.status_code,
                    response_data=error_data if 'error_data' in locals() else {}
                )
            
            return response
            
        except requests.exceptions.Timeout:
            raise NetworkException("请求超时")
        except requests.exceptions.ConnectionError:
            raise NetworkException("网络连接错误")
        except requests.exceptions.RequestException as e:
            raise NetworkException(f"网络请求失败: {str(e)}")
    
    def _validate_user_id(self, user_id: str):
        """验证用户ID格式"""
        if not user_id or not isinstance(user_id, str):
            raise ValidationException("UserID不能为空且必须是字符串")
        if len(user_id) < 1 or len(user_id) > 20:
            raise ValidationException("UserID长度必须在1-20个字符之间")
    
    # ==================== 会话管理 ====================
    
    def create_conversation(self, user_id: str, inputs: Dict[str, str] = None) -> str:
        """创建新的对话会话

        Args:
            user_id: 用户ID（1-20个字符）
            inputs: 变量输入（可选）

        Returns:
            str: 会话ID

        Raises:
            ValidationException: 参数验证错误
            APIException: API调用错误
        """
        self._validate_user_id(user_id)

        data = {
            "UserID": user_id,
            "Inputs": inputs or {}
        }

        # Debug logging
        import logging
        logger = logging.getLogger(__name__)
        logger.debug(f"create_conversation request data: {data}")

        response = self._make_request('POST', '/api/proxy/api/v1/create_conversation', data)

        # Log response for debugging 400 errors
        logger.debug(f"create_conversation response status: {response.status_code}")
        if response.status_code != 200:
            logger.error(f"create_conversation failed with status {response.status_code}")
            logger.error(f"Response text: {response.text}")
            logger.error(f"Request data was: {data}")

        result = response.json()

        if 'Conversation' not in result:
            raise APIException("API返回结果缺少'Conversation'字段")

        conversation_data = result['Conversation']
        if 'AppConversationID' not in conversation_data:
            raise APIException("API返回结果缺少'AppConversationID'字段")

        return Conversation.from_dict(conversation_data)
    
    def get_conversation_list(self, user_id: str, page: int = 1, 
                            page_size: int = 20) -> List[Conversation]:
        """获取用户的对话列表
        
        Args:
            user_id: 用户ID
            page: 页码（从1开始）
            page_size: 每页数量
            
        Returns:
            List[Conversation]: 对话列表
        """
        self._validate_user_id(user_id)
        
        data = {
            "UserID": user_id,
            "Page": page,
            "PageSize": page_size
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/get_conversation_list', data)
        result = response.json()
        
        conversations = []
        # 检查两种可能的字段名
        conv_list = result.get('ConversationList') or result.get('Conversations', [])
        for conv_data in conv_list:
            conversations.append(Conversation.from_dict(conv_data))
        
        return conversations
    
    def update_conversation(self, conversation_id: str, user_id: str, 
                          conversation_name: str) -> bool:
        """更新对话名称
        
        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            conversation_name: 新的会话名称
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id,
            "ConversationName": conversation_name
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/update_conversation', data)
        return response.ok
    
    def delete_conversation(self, conversation_id: str, user_id: str) -> bool:
        """删除对话
        
        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/delete_conversation', data)
        return response.ok
    
    def set_conversation_top(self, conversation_id: str, user_id: str) -> bool:
        """置顶会话
        
        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/set_conversation_top', data)
        return response.ok
    
    def cancel_conversation_top(self, conversation_id: str, user_id: str) -> bool:
        """取消会话置顶
        
        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/cancel_conversation_top', data)
        return response.ok
    
    # ==================== 聊天功能 ====================
    
    def chat(self, conversation_id: str, query: str, user_id: str,
             response_mode: str = "blocking", query_extends: QueryExtends = None) -> ChatResponse:
        """发送聊天消息（推荐使用，基于ChatQueryV2）

        Args:
            conversation_id: 会话ID
            query: 用户询问内容
            user_id: 用户ID
            response_mode: 响应模式，'streaming'或'blocking'，默认'blocking'
            query_extends: 查询扩展信息（如文件）

        Returns:
            ChatResponse: 聊天响应
        """
        self._validate_user_id(user_id)

        data = {
            "AppConversationID": conversation_id,
            "Query": query,
            "UserID": user_id,
            "ResponseMode": response_mode
        }

        if query_extends:
            data["QueryExtends"] = query_extends.to_dict()

        if response_mode == "streaming":
            return self._chat_streaming(data)
        else:
            return self._chat_blocking(data)
    
    def _chat_blocking(self, data: dict) -> ChatResponse:
        """阻塞式聊天"""
        response = self._make_request('POST', '/api/proxy/api/v1/chat_query_v2', data)
        result = response.json()
        
        return ChatResponse.from_dict(result)
    
    def _chat_streaming(self, data: dict) -> Iterator[dict]:
        """流式聊天，返回生成器

        Yields:
            dict: 直接返回HiAgent API的原始事件数据
        """
        response = self._make_request('POST', '/api/proxy/api/v1/chat_query_v2', data, stream=True)
        client = sseclient.SSEClient(response)

        for event in client.events():
            if event.data:
                try:
                    event_data = json.loads(event.data)
                    # 直接返回原始事件数据，不做任何处理
                    yield event_data
                except json.JSONDecodeError:
                    continue
    
    def chat_stream(self, conversation_id: str, query: str, user_id: str,
                   query_extends: QueryExtends = None) -> Iterator[dict]:
        """流式聊天（便捷方法）

        Args:
            conversation_id: 会话ID
            query: 用户询问内容
            user_id: 用户ID
            query_extends: 查询扩展信息

        Yields:
            dict: 包含 content 和 message_id 的字典
        """
        self._validate_user_id(user_id)

        data = {
            "AppConversationID": conversation_id,
            "Query": query,
            "UserID": user_id,
            "ResponseMode": "streaming"
        }

        if query_extends:
            data["QueryExtends"] = query_extends.to_dict()

        yield from self._chat_streaming(data)
    
    def chat_continue(self, conversation_id: str, user_id: str) -> ChatResponse:
        """继续对话
        
        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            
        Returns:
            ChatResponse: 聊天响应
        """
        self._validate_user_id(user_id)
        
        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id
        }
        
        return self._chat_blocking(data)
    
    def get_conversation_messages(self, conversation_id: str, user_id: str,
                                limit: int = 20) -> List[Message]:
        """获取会话历史消息列表
        
        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            limit: 限制返回条数（1-100）
            
        Returns:
            List[Message]: 消息列表
        """
        self._validate_user_id(user_id)
        
        if not 1 <= limit <= 100:
            raise ValidationException("limit参数必须在1-100之间")
        
        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id,
            "Limit": limit
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/get_conversation_messages', data)
        result = response.json()

        messages = []
        if 'Messages' in result:
            # 为了确保正确的时间排序，我们先收集所有消息并按时间戳排序
            all_messages = []

            for msg_data in result['Messages']:
                # 获取AI回复的时间戳作为对话时间
                answer_info = msg_data.get('AnswerInfo', {})
                created_timestamp = answer_info.get('CreatedTime', 0)

                # 添加用户消息（时间略早于AI回复）
                user_query = msg_data.get('Query', '')
                if user_query:
                    user_msg = Message(
                        message_id=msg_data.get('QueryID', ''),
                        conversation_id=msg_data.get('ConversationID', ''),
                        content=user_query,
                        role='user',
                        create_time=None,  # 用户消息时间通常不在这个API返回
                        feedback=None,
                        message_type='text',
                        status='completed'
                    )
                    # 用户消息排序时间戳 = AI回复时间戳 - 1秒
                    all_messages.append((created_timestamp - 1, user_msg))

                # 添加AI回复（使用from_dict处理）
                ai_msg = Message.from_dict(msg_data)
                if ai_msg.content:  # 只添加有内容的AI回复
                    all_messages.append((created_timestamp, ai_msg))

            # 按时间戳排序（正序：从早到晚）
            all_messages.sort(key=lambda x: x[0])

            # 提取排序后的消息
            messages = [msg for _, msg in all_messages]

        return messages
    
    def get_message_info(self, message_id: str, user_id: str) -> Message:
        """获取消息详情
        
        Args:
            message_id: 消息ID
            user_id: 用户ID
            
        Returns:
            Message: 消息详情
        """
        self._validate_user_id(user_id)
        
        data = {
            "MessageID": message_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/get_message_info', data)
        result = response.json()

        # 调试：打印实际返回的结果
        print(f"[DEBUG] get_message_info 原始响应: {result}")

        # 尝试处理不同的响应格式
        if 'MessageInfo' in result:
            # 真实 API 返回的是 MessageInfo 格式
            message_info = result['MessageInfo']
            answer_info = message_info.get('AnswerInfo', {})

            # 构造 Message 格式的数据
            message_data = {
                'AnswerInfo': answer_info,
                'Query': message_info.get('Query', ''),
                'ConversationID': message_info.get('ConversationID', '')
            }

            return Message.from_dict(message_data)
        elif 'Message' in result:
            # 兼容旧格式
            return Message.from_dict(result['Message'])
        elif isinstance(result, dict):
            # 如果是空结果，创建一个基本的Message对象
            if not any(result.values()) or all(v in [None, '', []] for v in result.values()):
                # 返回空Message但不抛出异常
                return Message(
                    message_id=message_id,
                    conversation_id="",
                    content="消息详情获取失败：HiAgent API返回空数据",
                    role="assistant",
                    create_time=None
                )
            # 尝试直接从结果创建Message
            return Message.from_dict(result)
        else:
            raise APIException(f"API返回结果格式不正确: {result}")
    
    def delete_message(self, message_id: str, user_id: str) -> bool:
        """删除消息
        
        Args:
            message_id: 消息ID
            user_id: 用户ID
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "MessageID": message_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/delete_message', data)
        return response.ok
    
    def feedback_message(self, message_id: str, user_id: str, feedback_type: str) -> bool:
        """消息反馈
        
        Args:
            message_id: 消息ID
            user_id: 用户ID
            feedback_type: 反馈类型，'like'或'dislike'
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        if feedback_type not in ['like', 'dislike']:
            raise ValidationException("feedback_type必须是'like'或'dislike'")
        
        data = {
            "MessageID": message_id,
            "UserID": user_id,
            "FeedbackType": feedback_type
        }

        # 调试：打印请求数据
        print(f"[DEBUG] feedback_message 请求数据: {data}")

        try:
            response = self._make_request('POST', '/api/proxy/api/v1/feedback', data)

            # 调试：打印响应
            try:
                result = response.json()
                print(f"[DEBUG] feedback_message 响应: {result}")
            except:
                print(f"[DEBUG] feedback_message 响应状态: {response.status_code}, 内容: {response.text}")

            return response.ok
        except APIException as e:
            if "404" in str(e):
                # HiAgent服务不可用时，返回False但不抛出异常
                print(f"[DEBUG] feedback_message HiAgent服务不可用: {e}")
                return False
            else:
                raise
    
    def stop_message(self, conversation_id: str, user_id: str) -> bool:
        """停止消息生成
        
        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/stop_message', data)
        return response.ok
    
    def get_suggested_questions(self, conversation_id: str, user_id: str,
                               message_id: str = None) -> List[str]:
        """获取提问建议

        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            message_id: 消息ID（可选，某些情况下必需）

        Returns:
            List[str]: 建议问题列表
        """
        self._validate_user_id(user_id)

        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id
        }

        # 如果提供了MessageID，添加到请求中
        if message_id:
            data["MessageID"] = message_id

        try:
            response = self._make_request('POST', '/api/proxy/api/v1/get_suggested_questions', data)
            result = response.json()
            return result.get('SuggestedQuestions', result.get('Questions', []))
        except APIException as e:
            # 如果因为缺少MessageID失败，并且还没有尝试过MessageID
            if "MessageID" in str(e) and not message_id:
                # 尝试获取最新的消息ID
                try:
                    messages = self.get_conversation_messages(conversation_id, user_id, limit=1)
                    if messages and len(messages) > 0:
                        latest_msg_id = messages[-1].message_id  # 获取最后一条消息的ID
                        if latest_msg_id:
                            return self.get_suggested_questions(conversation_id, user_id, latest_msg_id)
                except:
                    pass
            # 如果仍然失败，返回空列表
            return []
    
    def chat_query_v2(self, conversation_id: str, query: str, user_id: str,
                     response_mode: str = "blocking", query_extends: QueryExtends = None) -> ChatResponse:
        """ChatQueryV2接口（高级聊天功能）
        
        Args:
            conversation_id: 会话ID
            query: 用户询问内容
            user_id: 用户ID
            response_mode: 响应模式，'streaming'或'blocking'
            query_extends: 查询扩展信息
            
        Returns:
            ChatResponse: 聊天响应
        """
        self._validate_user_id(user_id)
        
        data = {
            "AppConversationID": conversation_id,
            "Query": query,
            "UserID": user_id,
            "ResponseMode": response_mode
        }
        
        if query_extends:
            data["QueryExtends"] = query_extends.to_dict()
        
        if response_mode == "streaming":
            return self._chat_streaming(data)
        else:
            return self._chat_blocking(data)
    
    def regenerate_response(self, conversation_id: str, user_id: str,
                          message_id: str) -> ChatResponse:
        """重新生成回复

        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            message_id: 消息ID（必需）

        Returns:
            ChatResponse: 新的聊天响应
        """
        self._validate_user_id(user_id)

        if not message_id:
            raise ValidationException("MessageID是必需参数")

        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id,
            "MessageID": message_id
        }

        response = self._make_request('POST', '/api/proxy/api/v1/query_again_v2', data)

        # 处理可能的空响应
        try:
            result = response.json()
        except:
            # 如果响应为空，创建一个默认响应
            result = {
                "MessageID": message_id,
                "Content": "",
                "ConversationID": conversation_id,
                "Status": "completed"
            }
        
        return ChatResponse.from_dict(result)
    
    def clear_messages(self, conversation_id: str, user_id: str) -> bool:
        """清除会话消息
        
        Args:
            conversation_id: 会话ID
            user_id: 用户ID
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "AppConversationID": conversation_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/clear_message', data)
        return response.ok
    
    
    
    # ==================== 工作流执行 ====================
    
    def run_workflow_async(self, user_id: str, inputs: Dict[str, Any] = None) -> str:
        """异步运行工作流
        
        Args:
            user_id: 用户ID
            inputs: 输入参数
            
        Returns:
            str: 任务ID
        """
        self._validate_user_id(user_id)
        
        data = {
            "UserID": user_id,
            "Inputs": inputs or {}
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/run_app_workflow', data)
        result = response.json()
        
        if 'TaskID' not in result:
            raise WorkflowException("工作流启动失败，缺少TaskID")
        
        return result['TaskID']
    
    def query_workflow_progress(self, task_id: str, user_id: str) -> WorkflowResult:
        """查询工作流执行进度
        
        Args:
            task_id: 任务ID
            user_id: 用户ID
            
        Returns:
            WorkflowResult: 工作流结果
        """
        self._validate_user_id(user_id)
        
        data = {
            "TaskID": task_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/query_run_app_process', data)
        result = response.json()
        
        return WorkflowResult.from_dict(result)
    
    def run_workflow_sync(self, user_id: str, inputs: Dict[str, Any] = None) -> WorkflowResult:
        """同步运行工作流
        
        Args:
            user_id: 用户ID
            inputs: 输入参数
            
        Returns:
            WorkflowResult: 工作流结果
        """
        self._validate_user_id(user_id)
        
        data = {
            "UserID": user_id,
            "Inputs": inputs or {}
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/sync_run_app_workflow', data)
        result = response.json()
        
        return WorkflowResult.from_dict(result)
    
    # ==================== 长期记忆管理 ====================
    
    def list_long_memory(self, user_id: str, page: int = 1, 
                        page_size: int = 20) -> List[LongMemory]:
        """获取长期记忆列表
        
        Args:
            user_id: 用户ID
            page: 页码
            page_size: 每页数量
            
        Returns:
            List[LongMemory]: 长期记忆列表
        """
        self._validate_user_id(user_id)
        
        data = {
            "UserID": user_id,
            "Page": page,
            "PageSize": page_size
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/list_long_memory', data)
        result = response.json()
        
        memories = []
        if 'Memories' in result:
            for memory_data in result['Memories']:
                memories.append(LongMemory.from_dict(memory_data))
        
        return memories
    
    def update_long_memory(self, memory_id: str, user_id: str, 
                          content: str) -> bool:
        """更新长期记忆
        
        Args:
            memory_id: 记忆ID
            user_id: 用户ID
            content: 记忆内容
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "MemoryID": memory_id,
            "UserID": user_id,
            "Content": content
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/update_long_memory', data)
        return response.ok
    
    def delete_long_memory(self, memory_id: str, user_id: str) -> bool:
        """删除长期记忆
        
        Args:
            memory_id: 记忆ID
            user_id: 用户ID
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "MemoryID": memory_id,
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/delete_long_memory', data)
        return response.ok
    
    def clear_long_memory(self, user_id: str) -> bool:
        """清除所有长期记忆
        
        Args:
            user_id: 用户ID
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/clear_long_memory', data)
        return response.ok
    
    # ==================== 用户变量管理 ====================
    
    def get_user_variables(self, user_id: str) -> List[UserVariable]:
        """获取用户环境变量
        
        Args:
            user_id: 用户ID
            
        Returns:
            List[UserVariable]: 用户变量列表
        """
        self._validate_user_id(user_id)
        
        data = {
            "UserID": user_id
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/get_app_user_variables', data)
        result = response.json()
        
        variables = []
        if 'Variables' in result:
            for var_data in result['Variables']:
                variables.append(UserVariable.from_dict(var_data))
        
        return variables
    
    def set_user_variables(self, user_id: str, 
                          variables: Dict[str, str]) -> bool:
        """设置用户环境变量
        
        Args:
            user_id: 用户ID
            variables: 变量字典
            
        Returns:
            bool: 是否成功
        """
        self._validate_user_id(user_id)
        
        data = {
            "UserID": user_id,
            "Variables": variables
        }
        
        response = self._make_request('POST', '/api/proxy/api/v1/set_app_user_variables', data)
        return response.ok